import logging
import re
from pathlib import Path

import numpy as np
from phonopy.api_phonopy import Phonopy
from phonopy.interface.vasp import read_vasp, read_vasp_from_strings

from .common_utils import ProcessingError, generate_pposcar_with_phonopy

logger = logging.getLogger(__name__)


def _extract_moment(line: str):
    _FLOAT = r"[-+]?(?:\d*\.\d+|\d+\.?(?:[eE][-+]?\d+)?)"
    RE_MOMENT = re.compile(
        rf"#\s*(?:cart\s*)?(?P<mx>{_FLOAT})\s+(?P<my>{_FLOAT})\s+(?P<mz>{_FLOAT})",
        re.IGNORECASE,
    )
    if "#" not in line:
        return (0.0, 0.0, 0.0)
    m = RE_MOMENT.search(line)
    return tuple(float(m.group(g)) for g in ("mx", "my", "mz")) if m else (0.0, 0.0, 0.0)


def _parse_poscar(poscar_content: str):
    lines = poscar_content.splitlines()

    match = re.search(r"OG\s*\(\s*(\d+)\.\s*(\d+)\.\s*(\d+)\s*\)", lines[0])
    if not match:
        raise ProcessingError("无法从POSCAR_msg的OG部分解析SG和OG编号")
    sg_number, og_number = int(match.group(1)), int(match.group(3))

    # --- 标题 & 缩放因子 ——
    scale = float(lines[1].split()[0])
    lattice = np.array([list(map(float, lines[i].split())) for i in range(2, 5)], float) * scale

    counts_idx = next(i for i, l in enumerate(lines) if re.fullmatch(r"\s*\d+(?:\s+\d+)+\s*", l))
    counts = list(map(int, lines[counts_idx].split()))
    nat = sum(counts)

    species_line = lines[counts_idx - 1].split()
    if all(re.fullmatch(r"\d+", s) for s in species_line):  # 老式 POSCAR 无名字
        species = [f"X{i+1}" for i in range(len(counts))]
    else:
        species = species_line

    coord_start = (
        next(
            i
            for i, l in enumerate(lines[counts_idx + 1 :], counts_idx + 1)
            if re.match(r"^(d|c)irect", l, re.IGNORECASE)
        )
        + 1
    )

    coord_lines = lines[coord_start : coord_start + nat]

    frac = np.zeros((nat, 3), float)
    mag = np.zeros((nat, 3), float)
    for idx, l in enumerate(coord_lines):
        frac[idx] = list(map(float, l.split()[:3]))
        mag[idx] = _extract_moment(l)

    return mag, frac, lattice, species, counts, sg_number, og_number


def _wrap_merge(frac: np.ndarray, mag: np.ndarray, latt: np.ndarray, tol_pos: float = 1e-3, tol_mag: float = 1e-3):
    """wrap 到 [0,1) 并去重，返回 (frac_u, mag_u, keep_idx)."""
    frac_wrapped = frac % 1.0
    cart = frac_wrapped @ latt

    keep: list[int] = []
    for i, (r_i, m_i) in enumerate(zip(cart, mag)):
        for j in keep:
            if np.linalg.norm(r_i - cart[j]) < tol_pos and np.linalg.norm(m_i - mag[j]) < tol_mag:
                break
        else:
            keep.append(i)
    idx = np.array(keep, int)
    return frac_wrapped[idx], mag[idx], idx


def _find_translation(frac1, frac2, tol=1e-5):
    N = frac1.shape[0]
    for i in range(N):
        for j in range(N):
            t = frac2[j] - frac1[i]
            t -= np.round(t)
            shifted = (frac1 + t) % 1.0
            for p in shifted:
                d2 = np.min(np.sum(((frac2 - p + 0.5) % 1.0 - 0.5) ** 2, axis=1))
                if d2 > tol**2:
                    break
            else:
                return t
    return None


def _write_poscar(filename, lattice, species, counts, frac, mag, comment: str = "Generated by wrap_merge"):
    with open(filename, "w") as f:
        f.write(comment + "\n")
        f.write("    1.0\n")
        for vec in lattice:
            f.write("  " + "  ".join(f"{x:16.9f}" for x in vec) + "\n")
        f.write("  " + "  ".join(f"{s:>4}" for s in species) + "\n")
        f.write("  " + "  ".join(f"{c:4d}" for c in counts) + "\n")
        f.write("Direct\n")
        for xyz, m in zip(frac, mag):
            f.write("  " + "  ".join(f"{v:12.8f}" for v in xyz))
            # 始终写磁矩（笛卡尔分量）
            f.write("  # cart " + "  ".join(f"{v:8.4f}" for v in m))
            f.write("\n")


async def run_mom2msg_workflow(poscar_msg_content: str):
    """
    执行从 mom2msg 借用的工作流来生成 POSCAR_clean。
    在指定的 cache_dir 中操作。
    """
    try:
        unitcell = read_vasp_from_strings(poscar_msg_content)
        ph = Phonopy(unitcell, primitive_matrix="auto")
        prim_matrix = ph.primitive_matrix
        if prim_matrix is None:
            raise ProcessingError("primitive_matrix is None")
        I = prim_matrix.T

        generate_pposcar_with_phonopy(Path("."), "POSCAR_msg")

        punitcell = read_vasp("PPOSCAR")
        frac_coords_pp = punitcell.scaled_positions
        R = (np.linalg.inv(I @ unitcell.cell)) @ punitcell.cell

        mag_cart, frac_coord, latt, species, counts, sg_number, og_number = _parse_poscar(poscar_msg_content)
        new_mag_cart = mag_cart @ R
        new_latt = I @ latt @ R
        new_frac_coord = frac_coord @ np.linalg.inv(I)

        frac_u, mag_u, keep_idx = _wrap_merge(new_frac_coord, new_mag_cart, new_latt)

        cum = np.cumsum([0] + counts)  # 元素块边界
        new_counts = [0] * len(counts)

        for i_orig in keep_idx:
            elem = next(k for k in range(len(counts)) if cum[k] <= i_orig < cum[k + 1])
            new_counts[elem] += 1

        translation = _find_translation(new_frac_coord, frac_coords_pp)
        if translation is not None:
            for i in range(len(frac_u)):
                frac_u[i] = (frac_u[i] + translation) % 1.0

        _write_poscar("POSCAR_clean", new_latt, species, new_counts, frac_u, mag_u, comment=f"generated by mom2msg")

        return sg_number, og_number, "  ".join(map(str, mag_u.flatten().round(4)))
    except Exception as e:
        raise ProcessingError(f"MOM2MSG工作流处理错误: {e}") from e
